import codecs

prompt_path = '/home/rickwwang/project_research/fairseq/data-bin/prompt.txt'
all_prompts = codecs.open(prompt_path, 'r', encoding='utf8').readlines()
all_prompts = map(lambda x: x.strip(), all_prompts)


# transformer-lm_bpe500_warmlr_format_pre2/27hyp.txt.1.topk40.t1.0
# transformer-lm_bpe500_warmlr_format2/40hyp.txt.1.topk40.t1.0
# bilstm_transformer-lm_bpe500_format_pre_srl_hier4_fix/25hyp.txt.own.prompt2srl2_500_wo_smooth_d1.1.topk10.t1.0
# bilstm_transformer-lm_bpe500_format_pre_srl_hier_pos_fix/17hyp.txt.own.prompt2srl2_500_wo_smooth_d1_verb_loss.1.topk10.t1.0
exp_file = []
# exp_file.append('fire_data/running/transformer-lm_bpe500_warmlr_format2/40hyp.txt.1.topk40.t1.0')
# exp_file.append('fire_data/running/transformer-lm_bpe500_warmlr_format_pre2/27hyp.txt.1.topk20.t1.0')
# exp_file.append('fire_data/running/bilstm_transformer-lm_bpe500_format_pre_srl_hier4_fix/25hyp.txt.own.prompt2srl2_500_wo_smooth_d1.1.topk20.t1.0')
exp_file.append('out/bilstm_transformer-lm_bpe500_format_pre_srl_hier_pos_fix/17hyp.txt.own.prompt2srl2_500_wo_smooth_d1_verb_loss.1.topk20.t1.0')


base_path = '/home/rickwwang/project_research/'

story_path = base_path + exp_file[0]
result_path = story_path + '.human'
all_stories = codecs.open(story_path, 'r', encoding='utf8').readlines()

promp_to_story = dict()
for s in all_stories:
    s = s.strip()
    key = s.split(' <SEP> ')[0]
    value = s.split(' <SEP> ')[1]
    if key in promp_to_story:
        promp_to_story[key].append(value)
    else:
        promp_to_story[key] = [value]

format_stories = []
for p in all_prompts:
    ss = promp_to_story[p]

    format_stories.append(p)
    format_stories.append(ss[0])

with codecs.open(result_path, 'w', encoding='utf8') as fout:
    fout.write('\n'.join(format_stories))
    fout.write('\n')